#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2; 

class KernelFastGelu {
public:
    __aicore__ inline KernelFastGelu() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t tilingKey)
    {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->blockLength = blockLength;
        this->tileNum =
            tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
        this->tileLength = tileLength / BUFFER_NUM;
        this->lasttileLength = lasttileLength;

        xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x + this->blockLength * GetBlockIdx(),
                            this->blockLength);
        yGm.SetGlobalBuffer((__gm__ DTYPE_X*)y + this->blockLength * GetBlockIdx(),
                            this->blockLength);

        pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(DTYPE_X));
        pipe.InitBuffer(tempBuf1, this->tileLength * sizeof(DTYPE_X));

        this->attr_half = 1.702 / 2;
    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        int32_t tileNum_ = this->tileNum;
        int32_t tileLength_ = this->tileLength;
        int32_t lasttileLength_ = this->lasttileLength;
        DTYPE_X attr_half_ = this->attr_half;

        for (int32_t i = 0; i < loopCount; i++) {
            CopyIn(i, tileNum_, tileLength_, lasttileLength_, loopCount);
            Compute(i, tileLength_, attr_half_);
            CopyOut(i, tileNum_, tileLength_, lasttileLength_, loopCount);
        }
    }

private:

    __aicore__ inline void CopyIn(int32_t progress, int32_t tileNum_, int32_t tileLength_, int32_t lasttileLength_, uint32_t loops)
    {
        LocalTensor<DTYPE_X> inLocal = inQueueIN.AllocTensor<DTYPE_X>();
        if (BUFFER_NUM == 2) {
            if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
                (progress == (this->tileNum * BUFFER_NUM - 1))) {
                DataCopy(
                    inLocal[0],
                    xGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                    (this->tileLength));
            }
            else {
                DataCopy(inLocal[0], xGm[progress * (tileLength_)],
                        (tileLength_));
            }
        }

        inQueueIN.EnQue(inLocal);
    }

    __aicore__ inline void Compute(int32_t progress, int32_t tileLength_, DTYPE_X attr_half_)
    {
        LocalTensor<DTYPE_X> inLocal = inQueueIN.DeQue<DTYPE_X>();
        LocalTensor<DTYPE_X> xLocal = inLocal;
        LocalTensor<DTYPE_Y> yLocal = outQueueOUT.AllocTensor<DTYPE_Y>();
        LocalTensor<DTYPE_X> temp1 = tempBuf1.Get<DTYPE_X>();

        Muls(temp1, xLocal, (DTYPE_X)-1.702, tileLength_);
        Exp(temp1, temp1, tileLength_);
        Adds(temp1, temp1, (DTYPE_X)1, tileLength_);
        Div(yLocal, xLocal, temp1, tileLength_);
        

        outQueueOUT.EnQue<DTYPE_Y>(yLocal);
        inQueueIN.FreeTensor(inLocal);
    }

    __aicore__ inline void CopyOut(int32_t progress, int32_t tileNum_, int32_t tileLength_, int32_t lasttileLength_, uint32_t loops)
    {
       LocalTensor<DTYPE_Y> outLocal = outQueueOUT.DeQue<DTYPE_Y>();
        if (BUFFER_NUM == 2) {
        if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
            (progress == (this->tileNum * BUFFER_NUM - 1))) {
            DataCopy(
                yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
                outLocal, (this->tileLength));
        }
        else {
            DataCopy(yGm[progress * (tileLength_)], outLocal,
                    (tileLength_));
        }
        }

        outQueueOUT.FreeTensor(outLocal);
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
    GlobalTensor<DTYPE_X> xGm;
    GlobalTensor<DTYPE_Y> yGm;
    TBuf<> tempBuf1;
    TBuf<> tempBuf2;
    TBuf<> tempBuf3;
    TBuf<> tempBuf4;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;
    uint32_t mask;
    uint32_t repeatTimes;
    DTYPE_X attr_half;
};


extern "C" __global__ __aicore__ void fast_gelu(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelFastGelu op;
    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }
    op.Init(x, y, 
          tiling_data.blockLength,
          tiling_data.tileNum, tiling_data.tileLength,
          tiling_data.lasttileLength, 
          tilingKey);
    op.Process();
}